import os.path

from langchain_community.document_loaders import Docx2txtLoader

from owl_admin.ext import db
from owl_ai.domain.chat_entity import ChatFileMessage, ChatMessage, Chat
from owl_ai.mapper import ChatFileMapper
from owl_ai.mapper.ai_config_mapper import AgentConfigMapper
from owl_ai.mapper.chat_message_mapper import ChatMessageMapper, ChatMapper
from owl_common.config import OWLConfig
from owl_common.sqlalchemy.transaction import Transactional
from .agent_flow_service import AgentFlowService, State
from .graph.agent_graph_compile import AgentGraphCompile
from ..domain.agent_config_entity import AgentConfigEntity


class ChatFilesService:

    @classmethod
    def upload_file(cls, chat_file_message: ChatFileMessage):
        """
        上传文件
        Returns:
        """
        return ChatFileMapper.insert_file(chat_file_message)

    @classmethod
    def find_file_by_id(cls, file_id: int):
        return ChatFileMapper.find_file_by_id(file_id)

    @classmethod
    def read_file_content(cls, file: ChatFileMessage):
        file_suffix = file.suffix
        file_path = os.path.join(OWLConfig().upload_path, file.path)
        if file_suffix == 'txt':
            return cls.read_txt_file(file_path)
        if file_suffix == 'pdf':
            return cls.read_pdf_file(file_path)
        if file_suffix == 'docx':
            return cls.read_docx_file(file_path)

    @classmethod
    def read_txt_file(cls, file_path):
        with open(file_path, 'r', encoding='utf-8') as f:
            content = f.read()
        return content

    @classmethod
    def read_pdf_file(cls, file_path):
        with open(file_path, 'rb') as f:
            content = f.read()
        return content

    @classmethod
    def read_docx_file(cls, file_path):
        loader = Docx2txtLoader(file_path)
        docs = loader.load()
        return docs[0].page_content


class ChatService:
    @classmethod
    @Transactional(db.session)
    def insert_chat(cls, chat: Chat):
        return ChatMapper.insert_chat(chat)

    @classmethod
    @Transactional(db.session)
    def find_chat_list(cls, chat: Chat):
        return ChatMapper.find_chat_list(chat)

    @classmethod
    @Transactional(db.session)
    def find_chat_by_id(cls, chat_id: int):
        return ChatMapper.find_chat_by_id(chat_id)


class ChatLLMService:

    @classmethod
    def chat_stream(cls, chat_id: int, chat_messages: list):
        # 获取agent配置
        # agent_id = chat.agent_id
        agent_config = AgentConfigMapper.get(AgentConfigEntity(id=1))

        # 组装agent对象
        # agent_flow = AgentFlowService.compile(agent=agent_config)
        agent_flow = AgentGraphCompile.compile(agent_config)

        state = State(messages=chat_messages, agent=agent_config)

        # 发起chat请求
        messages = agent_flow.stream(input=state, stream_mode="messages")
        return messages


class ChatMessageService:

    @classmethod
    @Transactional(db.session)
    def insert_message(cls, chat_message: ChatMessage):
        return ChatMessageMapper.insert_message(chat_message)

    @classmethod
    @Transactional(db.session)
    def find_messages_by_chat_id(cls, chat_id: int):
        return ChatMessageMapper.find_messages_by_chat_id(chat_id)
